{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "cell_0",
      "metadata": {},
      "source": [
        "# Advanced Usage Example - analyzing multiple documents with a single pipeline, with different LLMs, concurrency and cost tracking"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cell_1",
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install -U contextgem"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cell_2",
      "metadata": {},
      "source": [
        "To run the extraction, please provide your LLM details in the ``DocumentLLM(...)`` constructor further below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cell_3",
      "metadata": {},
      "outputs": [],
      "source": [
        "# Advanced Usage Example - analyzing multiple documents with a single pipeline,\n",
        "# with different LLMs, concurrency and cost tracking\n",
        "\n",
        "import os\n",
        "\n",
        "from contextgem import (\n",
        "    Aspect,\n",
        "    DateConcept,\n",
        "    Document,\n",
        "    DocumentLLM,\n",
        "    DocumentLLMGroup,\n",
        "    DocumentPipeline,\n",
        "    JsonObjectConcept,\n",
        "    JsonObjectExample,\n",
        "    LLMPricing,\n",
        "    NumericalConcept,\n",
        "    RatingConcept,\n",
        "    RatingScale,\n",
        "    StringConcept,\n",
        "    StringExample,\n",
        ")\n",
        "\n",
        "# Construct documents\n",
        "\n",
        "# Document 1 - Consultancy Agreement (shortened for brevity)\n",
        "doc1 = Document(\n",
        "    raw_text=(\n",
        "        \"Consultancy Agreement\\n\"\n",
        "        \"This agreement between Company A (Supplier) and Company B (Customer)...\\n\"\n",
        "        \"The term of the agreement is 1 year from the Effective Date...\\n\"\n",
        "        \"The Supplier shall provide consultancy services as described in Annex 2...\\n\"\n",
        "        \"The Customer shall pay the Supplier within 30 calendar days of receiving an invoice...\\n\"\n",
        "        \"All intellectual property created during the provision of services shall belong to the Customer...\\n\"\n",
        "        \"This agreement is governed by the laws of Norway...\\n\"\n",
        "        \"Annex 1: Data processing agreement...\\n\"\n",
        "        \"Annex 2: Statement of Work...\\n\"\n",
        "        \"Annex 3: Service Level Agreement...\\n\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "# Document 2 - Service Level Agreement (shortened for brevity)\n",
        "doc2 = Document(\n",
        "    raw_text=(\n",
        "        \"Service Level Agreement\\n\"\n",
        "        \"This agreement between TechCorp (Provider) and GlobalInc (Client)...\\n\"\n",
        "        \"The agreement shall commence on January 1, 2023 and continue for 2 years...\\n\"\n",
        "        \"The Provider shall deliver IT support services as outlined in Schedule A...\\n\"\n",
        "        \"The Client shall make monthly payments of $5,000 within 15 days of invoice receipt...\\n\"\n",
        "        \"The Provider guarantees [99.9%] uptime for all critical systems...\\n\"\n",
        "        \"Either party may terminate with 60 days written notice...\\n\"\n",
        "        \"This agreement is governed by the laws of California...\\n\"\n",
        "        \"Schedule A: Service Descriptions...\\n\"\n",
        "        \"Schedule B: Response Time Requirements...\\n\"\n",
        "    ),\n",
        ")\n",
        "\n",
        "# Create a reusable document pipeline for extraction\n",
        "contract_pipeline = DocumentPipeline()\n",
        "\n",
        "# Define aspects and aspect-level concepts in the pipeline\n",
        "# Concepts in the aspects will be extracted from the extracted aspect context\n",
        "contract_pipeline.aspects = [  # or use .add_aspects([...])\n",
        "    Aspect(\n",
        "        name=\"Contract Parties\",\n",
        "        description=\"Clauses defining the parties to the agreement\",\n",
        "        concepts=[  # define aspect-level concepts, if any\n",
        "            StringConcept(\n",
        "                name=\"Party names and roles\",\n",
        "                description=\"Names of all parties entering into the agreement and their roles\",\n",
        "                examples=[  # optional\n",
        "                    StringExample(\n",
        "                        content=\"X (Client)\",  # guidance regarding the expected output format\n",
        "                    )\n",
        "                ],\n",
        "            )\n",
        "        ],\n",
        "    ),\n",
        "    Aspect(\n",
        "        name=\"Term\",\n",
        "        description=\"Clauses defining the term of the agreement\",\n",
        "        concepts=[\n",
        "            NumericalConcept(\n",
        "                name=\"Contract term\",\n",
        "                description=\"The term of the agreement in years\",\n",
        "                numeric_type=\"int\",  # or \"float\", or \"any\" for auto-detection\n",
        "                add_references=True,  # extract references to the source text\n",
        "                reference_depth=\"paragraphs\",\n",
        "            )\n",
        "        ],\n",
        "    ),\n",
        "]\n",
        "\n",
        "# Define document-level concepts\n",
        "# Concepts in the document will be extracted from the whole document content\n",
        "contract_pipeline.concepts = [  # or use .add_concepts()\n",
        "    DateConcept(\n",
        "        name=\"Effective date\",\n",
        "        description=\"The effective date of the agreement\",\n",
        "    ),\n",
        "    StringConcept(\n",
        "        name=\"Contract type\",\n",
        "        description=\"The type of agreement\",\n",
        "        llm_role=\"reasoner_text\",  # for this concept, we use a more advanced LLM for reasoning\n",
        "    ),\n",
        "    StringConcept(\n",
        "        name=\"Governing law\",\n",
        "        description=\"The law that governs the agreement\",\n",
        "    ),\n",
        "    JsonObjectConcept(\n",
        "        name=\"Attachments\",\n",
        "        description=\"The titles and concise descriptions of the attachments to the agreement\",\n",
        "        structure={\"title\": str, \"description\": str | None},\n",
        "        examples=[  # optional\n",
        "            JsonObjectExample(  # guidance regarding the expected output format\n",
        "                content={\n",
        "                    \"title\": \"Appendix A\",\n",
        "                    \"description\": \"Code of conduct\",\n",
        "                }\n",
        "            ),\n",
        "        ],\n",
        "    ),\n",
        "    RatingConcept(\n",
        "        name=\"Duration adequacy\",\n",
        "        description=\"Contract duration adequacy considering the subject matter and best practices.\",\n",
        "        llm_role=\"reasoner_text\",  # for this concept, we use a more advanced LLM for reasoning\n",
        "        rating_scale=RatingScale(start=1, end=10),\n",
        "        add_justifications=True,  # add justifications for the rating\n",
        "        justification_depth=\"balanced\",  # provide a balanced justification\n",
        "        justification_max_sents=3,\n",
        "    ),\n",
        "]\n",
        "\n",
        "# Assign pipeline to the documents\n",
        "# You can re-use the same pipeline for multiple documents\n",
        "doc1.assign_pipeline(\n",
        "    contract_pipeline\n",
        ")  # assigns pipeline aspects and concepts to the document\n",
        "doc2.assign_pipeline(\n",
        "    contract_pipeline\n",
        ")  # assigns pipeline aspects and concepts to the document\n",
        "\n",
        "# Create an LLM group for data extraction and reasoning\n",
        "llm_extractor = DocumentLLM(\n",
        "    model=\"openai/gpt-4o-mini\",  # or any other LLM from e.g. Anthropic, etc.\n",
        "    api_key=os.environ[\"CONTEXTGEM_OPENAI_API_KEY\"],  # your API key\n",
        "    role=\"extractor_text\",  # signifies the LLM is used for data extraction tasks\n",
        "    pricing_details=LLMPricing(  # optional, for costs calculation\n",
        "        input_per_1m_tokens=0.150,\n",
        "        output_per_1m_tokens=0.600,\n",
        "    ),\n",
        ")\n",
        "llm_reasoner = DocumentLLM(\n",
        "    model=\"openai/o3-mini\",  # or any other LLM from e.g. Anthropic, etc.\n",
        "    api_key=os.environ[\"CONTEXTGEM_OPENAI_API_KEY\"],  # your API key\n",
        "    role=\"reasoner_text\",  # signifies the LLM is used for reasoning tasks\n",
        "    pricing_details=LLMPricing(  # optional, for costs calculation\n",
        "        input_per_1m_tokens=1.10,\n",
        "        output_per_1m_tokens=4.40,\n",
        "    ),\n",
        ")\n",
        "# The LLM group is used for all extraction tasks within the pipeline\n",
        "llm_group = DocumentLLMGroup(llms=[llm_extractor, llm_reasoner])\n",
        "\n",
        "# Extract all information from the documents at once\n",
        "doc1 = llm_group.extract_all(\n",
        "    doc1, use_concurrency=True\n",
        ")  # use concurrency to speed up extraction\n",
        "doc2 = llm_group.extract_all(\n",
        "    doc2, use_concurrency=True\n",
        ")  # use concurrency to speed up extraction\n",
        "# Or use async variants .extract_all_async(...)\n",
        "\n",
        "# Get the extracted data\n",
        "print(\"Some extracted data from doc 1:\")\n",
        "print(\"Contract Parties > Party names and roles:\")\n",
        "print(\n",
        "    doc1.get_aspect_by_name(\"Contract Parties\")\n",
        "    .get_concept_by_name(\"Party names and roles\")\n",
        "    .extracted_items\n",
        ")\n",
        "print(\"Attachments:\")\n",
        "print(doc1.get_concept_by_name(\"Attachments\").extracted_items)\n",
        "# ...\n",
        "\n",
        "print(\"\\nSome extracted data from doc 2:\")\n",
        "print(\"Term > Contract term:\")\n",
        "print(\n",
        "    doc2.get_aspect_by_name(\"Term\")\n",
        "    .get_concept_by_name(\"Contract term\")\n",
        "    .extracted_items[0]\n",
        "    .value\n",
        ")\n",
        "print(\"Duration adequacy:\")\n",
        "print(doc2.get_concept_by_name(\"Duration adequacy\").extracted_items[0].value)\n",
        "print(doc2.get_concept_by_name(\"Duration adequacy\").extracted_items[0].justification)\n",
        "# ...\n",
        "\n",
        "# Output processing costs (requires setting the pricing details for each LLM)\n",
        "print(\"\\nProcessing costs:\")\n",
        "print(llm_group.get_cost())\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}